from datetime import datetime
from itemadapter import ItemAdapter
from scrapy.exceptions import DropItem
from scrapy.utils.serialize import ScrapyJSONEncoder

from web_crawler.items import KlineItem
from web_crawler.utils import MySQLPool



class XueqiuKlinePipeline:
    """获取股票k线存入数据库"""

    def __init__(self):
        self.logger = None
        self.pool = MySQLPool()
        self.batch_size = 1000  # 批量插入条数
        self.batch_buffer = []  # 批量插入缓冲区

    def close_spider(self, spider):
        if self.batch_buffer:
            self._flush_batch()
        self.pool.close()
        self.logger.info("MySQL连接池已关闭")

    def open_spider(self, spider):
        """爬虫启动时自动调用"""
        self.logger = spider.logger
        self.create_table()
        spider.logger.info("K线Pipeline初始化完成")

    def create_table(self):
        """创建K线数据表（如果不存在）"""
        conn = None
        try:
            conn = self.pool.get_conn()
            with conn.cursor() as cursor:
                cursor.execute("""
                     CREATE TABLE IF NOT EXISTS stock_kline (
                        id BIGINT AUTO_INCREMENT PRIMARY KEY,
                        symbol VARCHAR(20) NOT NULL,
                        trade_date DATE NOT NULL,
                        open_price DECIMAL(12,4),
                        close_price DECIMAL(12,4),
                        high_price DECIMAL(12,4),
                        low_price DECIMAL(12,4),
                        volume BIGINT,
                        amount DECIMAL(20,4),
                        price_change DECIMAL(10,4),
                        percent_change DECIMAL(10,4),
                        turnover_rate DECIMAL(10,4),
                        macd_dea DECIMAL(12,4),
                        macd_dif DECIMAL(12,4),
                        macd DECIMAL(12,4),
                        boll_upper DECIMAL(12,4),
                        boll_lower DECIMAL(12,4),
                        ma20 DECIMAL(12,4),
                        kdj_k DECIMAL(12,4),
                        kdj_d DECIMAL(12,4),
                        kdj_j DECIMAL(12,4),
                        pe_ttm DECIMAL(12,4),
                        pb DECIMAL(12,4),
                        ps DECIMAL(12,4),
                        pcf DECIMAL(12,4),
                        market_cap DECIMAL(20,4),
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        UNIQUE KEY uk_symbol_date (symbol, trade_date),
                        INDEX idx_symbol (symbol),
                        INDEX idx_date (trade_date)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci
                    COMMENT='雪球K线数据表'
                """)
                conn.commit()
                self.logger.info("K线数据表验证/创建完成")
        except Exception as e:
            if conn:
                conn.rollback()
            self.logger.error(f"创建K线表失败: {e}")
            raise DropItem(f"数据库表创建失败: {e}")
        finally:
            if conn:
                conn.close()

    def process_item(self, item, spider):
        """处理并保存K线数据到MySQL"""
        # 检查Item类型
        if not isinstance(item, KlineItem):
            return item  # 如果不是KlineItem，直接返回
        # 数据验证
        adapter = ItemAdapter(item)
        required_fields = ['symbol', 'timestamp']
        for field in required_fields:
            if not adapter.get(field):
                raise DropItem(f"K线数据缺少必要字段({field})")

        trade_date = datetime.fromtimestamp(adapter['timestamp'] / 1000).date()
        # 添加到批量缓冲区
        self.batch_buffer.append({
            'symbol': adapter['symbol'],
            'trade_date': trade_date,  # 时间戳转日期
            'open_price': adapter.get('open_price'),
            'close_price': adapter.get('close_price'),
            'high_price': adapter.get('high_price'),
            'low_price': adapter.get('low_price'),
            'volume': adapter.get('volume'),
            'amount': adapter.get('amount'),
            'price_change': adapter.get('price_change'),
            'percent_change': adapter.get('percent_change'),
            'turnover_rate': adapter.get('turnover_rate'),
            'macd_dea': adapter.get('macd_dea'),
            'macd_dif': adapter.get('macd_dif'),
            'macd': adapter.get('macd'),
            'boll_upper': adapter.get('boll_upper'),
            'boll_lower': adapter.get('boll_lower'),
            'ma20': adapter.get('ma20'),
            'kdj_k': adapter.get('kdj_k'),
            'kdj_d': adapter.get('kdj_d'),
            'kdj_j': adapter.get('kdj_j'),
            'pe_ttm': adapter.get('pe_ttm'),
            'pb': adapter.get('pb'),
            'ps': adapter.get('ps'),
            'pcf': adapter.get('pcf'),
            'market_cap': adapter.get('market_cap')
        })

        # 达到批量大小时刷新
        if len(self.batch_buffer) >= self.batch_size:
            self._flush_batch()
        return item

    def _flush_batch(self):
        """执行批量插入"""
        if not self.batch_buffer:
            return

        conn = None
        try:
            conn = self.pool.get_conn()
            with conn.cursor() as cursor:
                # 构建批量插入SQL
                sql = """
                    INSERT INTO stock_kline (
                    symbol, trade_date, open_price, close_price,
                    high_price, low_price, volume, amount,
                    price_change, percent_change, turnover_rate,
                    macd_dea, macd_dif, macd, boll_upper, boll_lower,
                    ma20, kdj_k, kdj_d, kdj_j, pe_ttm, pb, ps, pcf, market_cap
                    ) VALUES (
                    %(symbol)s, %(trade_date)s, %(open_price)s, %(close_price)s,
                    %(high_price)s, %(low_price)s, %(volume)s, %(amount)s,
                    %(price_change)s, %(percent_change)s, %(turnover_rate)s,
                    %(macd_dea)s, %(macd_dif)s, %(macd)s, %(boll_upper)s, %(boll_lower)s,
                    %(ma20)s, %(kdj_k)s, %(kdj_d)s, %(kdj_j)s, %(pe_ttm)s, %(pb)s, %(ps)s, %(pcf)s, %(market_cap)s
                    )
                    ON DUPLICATE KEY UPDATE
                    open_price = VALUES(open_price),
                    close_price = VALUES(close_price),
                    high_price = VALUES(high_price),
                    low_price = VALUES(low_price),
                    volume = VALUES(volume),
                    amount = VALUES(amount),
                    updated_at = CURRENT_TIMESTAMP
                    """
                cursor.executemany(sql, self.batch_buffer)
                conn.commit()
                self.logger.info(f"批量插入 {len(self.batch_buffer)} 条K线数据")
        except Exception as e:
            if conn:
                conn.rollback()
            self.logger.error(f"批量插入失败: {e}")
        finally:
            if conn:
                conn.close()
            self.batch_buffer = []  # 清空缓冲区